package cn.iotab.iotkit.transport.mqtt.handler;

import cn.iotab.iotkit.transport.api.common.message.msg.rpc.AuthBasicRequestMsg;
import cn.iotab.iotkit.transport.api.common.message.msg.rpc.AuthResponseMsg;
import cn.iotab.iotkit.transport.api.common.message.msg.rpc.AuthX509RequestMsg;
import cn.iotab.iotkit.transport.api.common.message.msg.up.SessionEvent;
import cn.iotab.iotkit.transport.api.common.message.msg.up.SessionInfo;
import cn.iotab.iotkit.transport.api.common.queue.SendResultCallback;
import cn.iotab.iotkit.transport.api.common.util.EncryptionUtil;
import cn.iotab.iotkit.transport.api.common.util.SslUtil;
import cn.iotab.iotkit.transport.api.driver.DriverSession;
import cn.iotab.iotkit.transport.mqtt.MqttTransportContext;
import cn.iotab.iotkit.transport.mqtt.session.MqttSessionCtx;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.mqtt.MqttConnAckMessage;
import io.netty.handler.codec.mqtt.MqttConnAckVariableHeader;
import io.netty.handler.codec.mqtt.MqttConnectMessage;
import io.netty.handler.codec.mqtt.MqttConnectReturnCode;
import io.netty.handler.codec.mqtt.MqttFixedHeader;
import io.netty.handler.codec.mqtt.MqttMessageType;
import io.netty.handler.codec.mqtt.MqttQoS;
import io.netty.handler.ssl.SslHandler;
import lombok.extern.slf4j.Slf4j;

import javax.net.ssl.SSLPeerUnverifiedException;
import javax.security.cert.X509Certificate;
import java.net.InetSocketAddress;


@Slf4j
public class ConnectHandler extends AbstractMqttMessageHandler<MqttConnectMessage> {


    private final SslHandler sslHandler;


    private volatile InetSocketAddress address;

    public ConnectHandler(MqttTransportContext context, MqttSessionCtx deviceSessionCtx, SslHandler sslHandler) {
        super(context, deviceSessionCtx);
        this.sslHandler = sslHandler;
    }

    @Override
    public void process(ChannelHandlerContext ctx, MqttConnectMessage msg) {
        address = (InetSocketAddress) ctx.channel().remoteAddress();
        log.info("[{}] Processing connect msg for client: {}!", sessionId, msg.payload().clientIdentifier());
        X509Certificate cert;
        if (sslHandler != null && (cert = getX509Certificate()) != null) {
            processX509CertConnect(ctx, cert);
        } else {
            processAuthTokenConnect(ctx, msg);
        }

    }

    private X509Certificate getX509Certificate() {
        try {
            X509Certificate[] certChain = sslHandler.engine().getSession().getPeerCertificateChain();
            if (certChain.length > 0) {
                return certChain[0];
            }
        } catch (SSLPeerUnverifiedException e) {
            log.warn(e.getMessage());
            return null;
        }
        return null;
    }

    private void processX509CertConnect(ChannelHandlerContext ctx, X509Certificate cert) {
        try {
            if (!context.isSkipValidityCheckForClientCert()) {
                cert.checkValidity();
            }
            String strCert = SslUtil.getX509CertificateString(cert);
            String sha3Hash = EncryptionUtil.getSha3Hash(strCert);
            AuthX509RequestMsg authReq = AuthX509RequestMsg.builder().hash(sha3Hash).build();
            transportDriverService.authenticate(authReq,
                    new SendResultCallback<AuthResponseMsg>() {
                        @Override
                        public void onSuccess(AuthResponseMsg msg) {
                            onValidateDeviceResponse(msg, ctx);
                        }

                        @Override
                        public void onError(Throwable e) {
                            log.trace("[{}] Failed to process credentials: {}", address, sha3Hash, e);
                            ctx.writeAndFlush(createMqttConnAckMsg(MqttConnectReturnCode.CONNECTION_REFUSED_SERVER_UNAVAILABLE));
                            ctx.close();
                        }
                    });
        } catch (Exception e) {
            ctx.writeAndFlush(createMqttConnAckMsg(MqttConnectReturnCode.CONNECTION_REFUSED_NOT_AUTHORIZED));
            ctx.close();
        }
    }

    private void processAuthTokenConnect(ChannelHandlerContext ctx, MqttConnectMessage msg) {
        String userName = msg.payload().userName();
        log.info("[{}] Processing connect msg for client with user name: {}!", sessionId, userName);
        AuthBasicRequestMsg.AuthBasicRequestMsgBuilder authReqBuilder = AuthBasicRequestMsg.builder();
        authReqBuilder.clientId(msg.payload().clientIdentifier());
        if (userName != null) {
            authReqBuilder.userName(userName);
        }
        String password = msg.payload().password();
        if (password != null) {
            authReqBuilder.password(password);
        }
        AuthBasicRequestMsg authBasicRequestMsg = authReqBuilder.build();
        transportDriverService.authenticate(authBasicRequestMsg,
                new SendResultCallback<AuthResponseMsg>() {
                    @Override
                    public void onSuccess(AuthResponseMsg msg) {
                        onValidateDeviceResponse(msg, ctx);
                    }

                    @Override
                    public void onError(Throwable e) {
                        log.info("[{}] Failed to process credentials: {}", address, userName, e);
                        ctx.writeAndFlush(createMqttConnAckMsg(MqttConnectReturnCode.CONNECTION_REFUSED_SERVER_UNAVAILABLE));
                        ctx.close();
                    }
                });
    }

    private void onValidateDeviceResponse(AuthResponseMsg msg, ChannelHandlerContext ctx) {
        if (!msg.isSuccess()) {
            ctx.writeAndFlush(createMqttConnAckMsg(MqttConnectReturnCode.CONNECTION_REFUSED_NOT_AUTHORIZED));
            ctx.close();
        } else {
            //认证成功, 补充SessionInfo信息
            deviceSessionCtx.setConnected();
            //构造SessionInfo
            SessionInfo sessionInfo = SessionInfo.of(deviceSessionCtx.getSessionId(), msg.getProductKey(), msg.getDeviceId());

            //上报Session状态
            transportDriverService.onSessionStateChange(sessionInfo, SessionEvent.OPEN, new SendResultCallback<Void>() {
                @Override
                public void onSuccess(Void msg) {
                    SessionMsgHandler sessionMsgHandler = new SessionMsgHandler(deviceSessionCtx);
                    DriverSession driverSession = transportDriverService.registerSession(sessionInfo, sessionMsgHandler);
                    deviceSessionCtx.setDriverSession(driverSession);
                    ctx.writeAndFlush(createMqttConnAckMsg(MqttConnectReturnCode.CONNECTION_ACCEPTED));
                    log.info("[{}] Client {}:{} connected!", sessionId, sessionInfo.getProductKey(), sessionInfo.getDeviceId());
                }

                @Override
                public void onError(Throwable e) {
                    log.warn("[{}] Failed to submit session event", sessionId, e);
                    ctx.writeAndFlush(createMqttConnAckMsg(MqttConnectReturnCode.CONNECTION_REFUSED_SERVER_UNAVAILABLE));
                    ctx.close();
                }
            });
        }
    }

    private MqttConnAckMessage createMqttConnAckMsg(MqttConnectReturnCode returnCode) {
        MqttFixedHeader mqttFixedHeader =
                new MqttFixedHeader(MqttMessageType.CONNACK, false, MqttQoS.AT_MOST_ONCE, false, 0);
        MqttConnAckVariableHeader mqttConnAckVariableHeader = new MqttConnAckVariableHeader(returnCode, true);
        return new MqttConnAckMessage(mqttFixedHeader, mqttConnAckVariableHeader);
    }
}
